"""Particle filter"""

import numpy as np
from copy import deepcopy

class Particle(object):
    """
    Particle representation
    """

    def __init__(self, model, weight, params=None):
        self._weight = weight
        self._model = model
        self._params = params


class Particle_filter(object):
    """
    Particle filter

    Attributes
    ----------
    size : int
        Number of particles
    particles : list
        Array of particles
    resampling_threshold : int
        Threshold for resampling
    _proposal : some pdf
        Proposal pdf
    _no_of_resamplings : int
        Number of resamplings (counter)

    Methods
    -------
    __init__(model, size, proposal, resampling_threshold_ratio=.5)
        Constructor
    populate_particles()
        Method populating particles for the filter
    update(data)
        Update particles
    particles_data_update(data)
        Data update of particles
    particles_time_update()
        Time update of particles
    normalize_weights()
        Normalizes weights of particles
    resample()
        Performs resampling
    predict(regression_vector)
        Perform prediction (as weighted sum)
    log()
        Write log
    """

    def __init__(self, model, size, proposal, resampling_threshold_ratio=.5):
        """Initialize particle filter"""
        self.size = size
        self.particles = []
        self._proposal = proposal
        self.populate_particles()
        self.resampling_threshold = np.floor(size * resampling_threshold_ratio)
        self._no_of_resamplings = 0
        
    def populate_particles(self):
        """Populates particles"""
        #particle = Particle(model, 1./float(self.size), self._proposal.sample())
        init_weight = 1./float(self.size)
        for i in xrange(self.size):
       	    self.particles.append(Particle(model, init_weight, self._proposal.sample()))
       	    #deepcopy(particle))
       	    #self.particles.append(deepcopy(particle))

    def update(self, data):
        """Update particles and resample if necessary"""
        self.particles_time_update()
        self.particles_data_update(data)
        self.normalize_weights()
       	weights = np.array([particle._weight for particle in self.particles])
       	params = np.array([particle._params for particle in self.particles]).T

        log = np.vstack((weights, params))
        log = np.vstack((log, np.arange(self.size)))
        np.set_printoptions(precision=5, suppress=True)
        print "\n{0:6s}".format(log.T)

       	N_eff = 1 / np.sum([particle._weight ** 2 for particle in self.particles])
       	if N_eff < self.resampling_threshold:
            res = self.resample()
            indices = [i for i in res]
            print indices[0]
            self._no_of_resamplings += 1


    def particles_data_update(self, data):
        for particle in self.particles:
           	particle._model.model_update(data, particle._params)
           	particle._weight *= np.exp(particle._model.delta_log_likelihood)

    def particles_time_update(self):
        for particle in self.particles:
       	    particle._params = self._proposal.sample(particle._params)

    def normalize_weights(self):
        """Normalize particles' weights"""
        weights_sum = np.sum(self.particles[i]._weight for i in xrange(self.size))
        for particle in self.particles:
            particle._weight /= weights_sum 

    def resample(self):
        weights = [particle._weight for particle in self.particles]
        n = len(weights)
        cum_weights = np.cumsum(weights)
        rand_draw = np.random.random()
        rand_draws = [(rand_draw + i ) / n for i in xrange(n)]
        indices = np.searchsorted(cum_weights, rand_draws)
        yield np.array(indices)
        tmp = []
        for i in indices:
        	self.particles[i]._weight = 1./n
        	tmp.append(deepcopy(self.particles[i]))
        self.particles = tmp


    def predict(self, regression_vector):
        """Prediction - is it really necessary?"""
        pass

    def log(self):
        print "Resamplings: %d" % self._no_of_resamplings



class Dirichlet(object):
    """
    Dirichlet class for particle filtering random walk

    Attributes
    ----------
    random_walk_width : float
        Width of the random walk
    init_params : float
        Initial parameters for the distribution
    stabilization : float
        Walk stabilizing term

    Methods
    -------
    sample(params=None)
        Sampler method
    """

    def __init__(self, init_params, random_walk_width=0.001, stabil=0.01):
        """Constructor"""
        self.random_walk_width = random_walk_width
        self.init_params = init_params
        self.stabilization = stabil

    def sample(self, params=None):
        """
        Sampler

        If params == None, uses the initial parameters,
        """

        if params == None:
        	params = self.init_params

        if type(params).__name__ == 'float' or \
           type(params).__name__ == 'numpy.float64':
            params = np.array([params, 1. - params]).T
            return np.random.dirichlet(params/self.random_walk_width
                   + self.stabilization)[0]
        else:
            params = np.array(params).T
            return np.random.dirichlet(params/self.random_walk_width
                   + self.stabilization)
       
        

# Dirichlet rand. walk pro vyvoj parametru zapominani
# weights = np.array([.4, .5, .1])
# weights.new = np.random.dirichlet(weights/.001)
#===========================
if __name__ == "__main__":
    from giw import *
    print "Test"

    data_real = [7.]
    theta = (0.9, -0.5)
    for i in xrange(1, 100):
        data_real.append(theta[0] * data_real[i-1] \
                         + theta[1] + np.random.randn(1) * 1. + i/10)
    datamatrix = data_matrix(data_real, 1)
    data = iter(datamatrix)

    from regmodelpf import *
    V = np.eye(3) * 0.01
    V[0,0] = 0.1
    model = RegmodelPF(V)

    proposal = Dirichlet([.95, .01, .01, .3], 0.001, 0.001)

    size = 55
    pf = Particle_filter(model, size, proposal, 0.1)
#    for i in xrange(number):
#    	pf.particles[i]._params = 0.9 + float(i) * 0.01
#    raw_input()
    for dt in data:
    	pf.update(dt)
    pf.log()
